import re

from gensim.models import Word2Vec
import pymongo
from matplotlib.font_manager import FontProperties
import jieba

host = '192.168.0.250'
port = 27018
username = 'bashou_dev'
password = 'bashoukejihhrhl123'

mongo_client = pymongo.MongoClient(host, port, connect=False, socketKeepAlive=True, wtimeout=10000, w=1,
                                   waitQueueTimeoutMS=10000)
mongo_client.admin.authenticate(username, password)
db = mongo_client['develop']

stoplist = {}.fromkeys([line.strip() for line in open("./stopwords.txt")])

PAGE_SIZE = 1000
MAX_PAGE = 100
begin_id = '0'
model = Word2Vec.load('./word2vec.model')

for i in range(0, MAX_PAGE):
    print('current_page:' + str(i))
    instruments = db.instrument.find({'_id': {'$gt': begin_id}}, {'_id': 1, 'caseTitle': 1, 'docContent': 1}) \
        .sort("_id", 1).limit(PAGE_SIZE)
    docs = [x for x in instruments]
    begin_id = docs[-1]['_id']
    print("loaded data------------")

    sentences = []
    for doc in docs:
        for sentence in doc['docContent'].split("。"):
            sentences.append([word for word in (jieba.lcut(re.sub('[\d ×()]', '', sentence))) if word not in stoplist])
    print("cut word finish------------")

    # print(sentences)
    # if i == 0:
    #     model = Word2Vec(sentences, sg=1, hs=1, min_count=1, window=3, size=200, workers=4)
    # else:
    model.train(sentences, epochs=model.epochs, total_examples=model.corpus_count)
    print("model finish------------")

model.save('./word2vec.model')

# req_count = 10
# for key in model.wv.similar_by_word('广东省', topn=10):
#     req_count -= 1
#     print(key[0], key[1])
#
# print(model.wv.similarity('广东省', '珠海市'))
# print(model.wv.similarity('珠海市', '西安市'))
# print(model.wv.doesnt_match("广东省 珠海市 西安市".split()))

# from sklearn.manifold import TSNE
#
# from itertools import chain
#
# random_word = list(set(list(chain.from_iterable(sentences))))[0:300]
#
# X_tsne = TSNE(n_components=2, init='pca', random_state=501).fit_transform(model.wv[random_word])
#
# import matplotlib.pyplot as plt
#
# # 解决负号'-'显示为方块的问题
# plt.figure(figsize=(14, 8))
# # myfont = FontProperties(fname='/usr/share/fonts/wqy-zenhei/wqy-zenhei.ttc')
# myfont = FontProperties(fname='/Library/Fonts/Songti.ttc')
#
# plt.scatter(X_tsne[:, 0], X_tsne[:, 1])
# for i in range(len(X_tsne)):
#     x = X_tsne[i][0]
#     y = X_tsne[i][1]
#     plt.text(x, y, random_word[i], fontproperties=myfont, size=10)
#
# plt.show()
# print('finish')
